{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# step1 导入相关库\n",
    "# 使用的是 pytorch 框架\n",
    "import copy\n",
    "import json\n",
    "import torch\n",
    "# import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from tqdm import tqdm\n",
    "from typing import List\n",
    "from einops import rearrange\n",
    "from datasets import load_dataset\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import AutoConfig, AutoTokenizer, AutoModel, AutoModelForCausalLM\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建 LoraLinear 类\n",
    "class LoraLinear(torch.nn.Module):\n",
    "  def __init__(\n",
    "    self,\n",
    "    base_layer: torch.nn.Linear, # 需要被替换的线性层，也就是原来的线性层\n",
    "    r: int = 8, # Lora 的秩 也就是  lora rank\n",
    "    lora_alpha: int = 16, # Lora 的缩放因子\n",
    "    lora_dropout: float = 0.05, # Lora 的 dropout 概率\n",
    "    test_mode: bool = False, # 是否是测试模式, 用于控制 lora_B 是否为全零\n",
    "  ):\n",
    "    super(LoraLinear, self).__init__()\n",
    "    self.base_layer = base_layer\n",
    "    self.r = r\n",
    "    self.lora_alpha = lora_alpha\n",
    "    self.dropout = torch.nn.Dropout(p=lora_dropout)\n",
    "\n",
    "    # 初始化 lora_A 和 lora_B\n",
    "    self.lora_A = torch.nn.Parameter(torch.randn(r, base_layer.in_features) * 0.01)\n",
    "    self.lora_B = torch.nn.Parameter(torch.randn(base_layer.out_features, r))\n",
    "\n",
    "    # 有了 lora_A  和  lora_B 我们就来初始化 lora 矩阵\n",
    "    # 如果是测试模式, 则将 lora_B 初始化为全零\n",
    "\n",
    "    torch.nn.init.normal_(self.lora_A, mean=0.0, std=0.01)\n",
    "    if test_mode:\n",
    "      torch.nn.init.zeros_(self.lora_B)\n",
    "    else:\n",
    "      torch.nn.init.normal_(self.lora_B, mean=0.0, std=0.01)\n",
    "    \n",
    "    # 冻结原来的层的参数(原来的线性层)\n",
    "    for param in self.base_layer.parameters():\n",
    "      param.requires_grad = False\n",
    "\n",
    "  # 实现前向传播, 也就是forward函数\n",
    "  def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "    # 实际版本是正确的：\n",
    "    # x → A → B → scaling\n",
    "    # lora_adjustment = F.linear(x, self.lora_A)        # x × A^T\n",
    "    # lora_adjustment = F.linear(lora_adjustment, self.lora_B)  # (x × A^T) × B^T\n",
    "    # lora_adjustment = lora_adjustment * scaling\n",
    "\n",
    "    scaling = float(self.lora_alpha) / float(self.r)     # lora 缩放系数\n",
    "    # lora_adjustment = F.linear(self.dropout(x), self.lora_A) x * A^T\n",
    "    lora_adjustment = F.linear(F.linear(self.dropout(x), self.lora_A), self.lora_B)\n",
    "    return self.base_layer(x) + lora_adjustment * scaling"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lora_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.18 | packaged by Anaconda, Inc. | (main, Jun  5 2025, 13:08:55) [MSC v.1929 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "6e13b17ab0ad3c103e5da93217e1b059758434f0f7c11cc7a16daec7a779cc3c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
